#! /usr/bin/env python
'''
GUI-less interactive script (RotaPy Jr.)

Specify a set of 3-component input seismograms and a 'correct' backazimuth
and calculate the inverse-weighted ratio of the envelope functions of 
radial and transverse components of the seismograms.

Maximum envelope ratio (radial/transverse) should occur at a backazimuth
that is correct. Orientation correct for sensors is then calculated as the
difference between the starting backazimuth and the calculated backazimuth.
'''
import os.path

import numpy as np
import matplotlib.pyplot as plt

from obspy.core import read
from obspy.signal import filter
from obspy.signal import rotate
from matplotlib.widgets import Slider
from matplotlib.widgets import Button
from matplotlib.widgets import MultiCursor
from matplotlib.widgets import SpanSelector

__VERSION__ = '0.1'

def update(val):
	"""
	When slider value is updated, recalculate the time plot traces
	and replot the radial/transverse components for the new backazimuth
	"""
	azi = slider_azimuth.val
	radtmp, transtmp = rotate.rotate_NE_RT(trace_N,trace_E,azi)
	line_rad.set_ydata(radtmp)
	line_trans.set_ydata(transtmp)
	line_rad_env.set_ydata(filter.envelope(radtmp))
	line_trans_env.set_ydata(filter.envelope(transtmp))
	
	fig.canvas.draw()
	
def reset(event):
	"""
	Reset the azimuth slider on button-press
	"""
	slider_azimuth.reset()

def onSelect(xmin, xmax):
	"""
	Obtain indices of selected window (in time plot) and 
	calculate the envelope and ratios over the range of 0-360
	degrees backazimuth
	"""
	indmin, indmax = np.searchsorted(time, (xmin, xmax))
	indmax = min(len(time)-1, indmax)
	
	# Obtain the maximum radial, transverse, and envelope values
	# over the range of 0-360 degrees over window
	max_radial, max_transverse, max_env = bazSweep(indmin,indmax) 
	
	# Reset and re-plot using windowed data in bottom-left plot
	ax2.cla()
	ax2.plot(azimuths,max_radial,'r',azimuths,max_transverse,'b')
	ax2.set_xlabel("Azimuth [deg]")
	ax2.set_ylabel("Maxium Envelope")
	ax2.set_xlim([0,360])
	# Reset and re-plot using windowed data in bottom-right plot
	ax3.cla()
	ax3.plot(azimuths,max_env,'g')
	ax3.set_xlim([0,360])
	ax3.set_xlabel("Azimuth[deg]")
	ax3.set_ylabel("Envelope Ratio (R/T)")

	fig.canvas.draw()

def bazSweep(indmin,indmax):
	"""
	Sweep over 0-360 degrees, and obtain maximum R/T/E values
	and return
	"""
	# Initialize empty lists to grow
	max_radial = []
	max_transverse = []
	max_env = []
	azimuths = np.arange(MIN_ANG,MAX_ANG) # 0-360
	
	# For every angle in range of 0 -> 360
	for theta in azimuths:
		# Calculate new radial/transverse components using theta as backazimuth
		r, t = rotate.rotate_NE_RT(trace_N,trace_E,theta)
		rad_tmp_env = filter.envelope(r) # Calculate envelope of radial
		trans_tmp_env = filter.envelope(t) # Calculate envelope of transverse
		rad_tmp_env = rad_tmp_env[indmin:indmax] # Extract the envelope between the xmin:xmax indicies
		trans_tmp_env = trans_tmp_env[indmin:indmax] # Extract the envelope between the xmin:xmax indicies
		
		max_radial.append(np.max(rad_tmp_env)) # Obtain maximum envelope for radial at this backazimuth
		max_transverse.append(np.max(trans_tmp_env)) # Obtain maximum envelope for transverse at this backazimuth
		
		# Calculate ratio of radial/transverse components 
		env_ratio = np.true_divide(rad_tmp_env,trans_tmp_env)
		# Weight envelope ratio by the inverse amplitude of radial component
		weight_matrix = np.true_divide(1.,rad_tmp_env)
		weight_ratio = np.multiply(weight_matrix,env_ratio)
		# Find maximum weighted-envelope value and add to list
		max_env.append(np.max(weight_ratio))

	return max_radial, max_transverse, max_env
	

if __name__ == "__main__":

	MIN_ANG = 0 # Minimum angle for azimuth sweep
	MAX_ANG = 359 # Maximum angle for azimuth sweep
	ANGLE = 155 # Alleged back-azimuth for a particular sensor/line location
	inputdir = r'C:\Users\Adam\Desktop\geometry\eq' # Directory with data files
	#inputdir = r'C:\Users\Adam\Desktop\geometry\data' # Directory with data files
	
	
	"""
	# A few set of stations we're testing the  method on here.
	# Using earthquakes appears to give a more stable result due to longer wavelength and 
	# the delay time between P-arrivals/rest of wavetrain compared to the active-source
	# data.
	# 
	# These files are listed below. Comment/uncomment as needed
	"""
	# GeoNet station, BETTER BE CORRECT! - Angle: 155
	z = 'NZ.OXZ.Z.SAC'
	n = 'NZ.OXZ.N.SAC'
	e = 'NZ.OXZ.E.SAC'
	
	# Believed to be 'correct' - Angle: 139
	#z = 'S002.PEG21.1388.1390.Z.sgy'
	#n = 'S002.PEG21.1388.1390.N.sgy'
	#e = 'S002.PEG21.1388.1390.E.sgy'
	
	# Believed to be 'wrong' by 22 degrees. Angle: 147
	#z = 'S028.PEG17.1470.1470.Z.sgy'
	#n = 'S028.PEG17.1470.1470.N.sgy'
	#e = 'S028.PEG17.1470.1470.E.sgy'
	
	# Believed to be 'wrong' by unknown amount (cmpaz says 44 degrees?). Angle: 154
	#z = 'S001.PEG19.1538.1554.Z.sgy'
	#n = 'S001.PEG19.1538.1554.N.sgy'
	#e = 'S001.PEG19.1538.1554.E.sgy'
	
	
	# Read in data for Z,N,E traces and stack
	for i, tr in enumerate(read(os.path.join(inputdir,z))):
		if i == 0:
			delta = tr.stats.delta
			npts = tr.stats.npts
			time = np.arange(0, delta*(npts-0.5), delta)
			trace_Z = tr.data
		else:
			trace_Z += tr.data
	for i, tr in enumerate(read(os.path.join(inputdir,n))):
		if i == 0:
			trace_N = tr.data
		else:
			trace_N += tr.data
			
	for i, tr in enumerate(read(os.path.join(inputdir,e))):
		if i == 0:
			trace_E = tr.data
		else:
			trace_E += tr.data
	
	# Initialize azimuth vector, and empty arrays for growing the maximum envelopes
	azimuths = np.arange(MIN_ANG,MAX_ANG)
#	max_transverse = []
#	max_radial = []
#	
	# Calculate the initial radial/transverse components with starting backazimuth 
	rad, trans = rotate.rotate_NE_RT(trace_N,trace_E,ANGLE)
	radial_env = filter.envelope(rad) # Calculate the envelope of radial component
	transverse_env = filter.envelope(trans) # Calculate the envelope of transverse component

	# Setup the plot and add the interactive widgets
	
	fig = plt.figure()
	
	ax1 = plt.subplot2grid((6,6),(0,0),colspan=6,rowspan=3)
	ax2 = plt.subplot2grid((6,6),(3,0),colspan=3,rowspan=3)
	ax3 = plt.subplot2grid((6,6),(3,3),colspan=3,rowspan=3,sharex=ax2)
	
	# Label the axes
	ax1.set_xlabel("Time [s]")
	ax1.set_ylabel("Amplitude")
	ax2.set_xlabel("Azimuth [deg]")
	ax2.set_ylabel("Maxium Envelope")
	ax2.set_xlim([0,360])
	ax3.set_xlabel("Azimuth [deg]")
	title = "%s.%s" % (z.split('.')[0],z.split('.')[1])
	ax1.set_title(title)
	
	# Plot the initial set of data (envelopes and traces)
	line_vert, line_rad, line_rad_env, line_trans, line_trans_env, = \
		ax1.plot(time,trace_Z,'k',time,rad,'r',time,radial_env,'r:',
				time,trans,'b',time,transverse_env,'b:')
	
	# Create widgets (span selector)
	span = SpanSelector(ax1, onSelect, 'horizontal', useblit=True,
				rectprops=dict(alpha=0.5, facecolor='#00FF00') )
	# Create azimuth slider
	axes_azimuth = plt.axes([0.1,0.03,0.65,0.03])
	slider_azimuth = Slider(axes_azimuth,'Azimuth',0,359,valinit=ANGLE)
	slider_azimuth.on_changed(update)
	# Create azimuth slider reset button
	axes_reset_button = plt.axes([0.8, 0.025, 0.1, 0.04])
	button = Button(axes_reset_button, 'Reset',color='r')
	button.on_clicked(reset)
	# Create linked-cursor between the two bottom plots
	cursor = MultiCursor(fig.canvas,(ax2,ax3),color='r',lw=1)
	# Adjust the borders of the figure to fit our setup
	fig.subplots_adjust(hspace=0.75,wspace=0.70,left=0.08,right=0.98,top=0.95,bottom=0.13)
	# Show the figure
	plt.show()